import ctypes
import time

import matplotlib.pyplot as plt
import numpy as np

from device import AC_SETFUNCTION_BASELINECLAMP, DRV_IDLE, Device

# conf

dll_file_path = "atmcd64d.dll"
img_file_path = "image.bin"

#  run

with Device(dll_file_path) as device:
    # %%
    # fetch info

    capabilities = device.get_capabilities()
    print(f"got {capabilities=}")

    head_model_name = device.get_head_model()
    print(f"got {head_model_name=}")

    xn, yn = device.get_detector_size()
    print(f"got detector_size: {xn=}, {yn=}")

    # %%
    # set parameters

    acquisition_mode = 1
    device.set_acquisition_mode(acquisition_mode)  # Single scan
    print(f"set {acquisition_mode=}")

    exposure_time = 0.001
    device.set_exposure_time(exposure_time)
    print(f"set {exposure_time=}")

    exposure_time, accumulate_time, kinetic_time = device.get_acquisition_times()
    print(f"got {exposure_time=}, {accumulate_time=}, {kinetic_time=}")

    read_mode = 4
    device.set_read_mode(read_mode)
    print(f"set read_mode: { read_mode=}")

    vs_speed_index, vs_speed = device.get_fastest_recommended_vs_speed()
    print(f"got (recommended) {vs_speed_index=}, {vs_speed=}")

    device.set_vs_speed(vs_speed_index)
    print(f"set {vs_speed_index=}")

    ad_channel_num = device.get_ad_channels_num()
    print(f"got {ad_channel_num=}")

    typ = 0
    best_ad_channel_index, best_hs_speed_index, best_hs_speed = 0, 0, 0
    for ad_channel_index in range(ad_channel_num):
        print(f"with ad_channel_index = {ad_channel_index}")

        hs_speed_num = device.get_hs_speeds_num(ad_channel_index, typ)
        print(f"\tgot {hs_speed_num=}")

        for hs_speed_index in range(hs_speed_num):
            print(f"\twith hs_speed_index = {hs_speed_index}")

            hs_speed = device.get_hs_speed(ad_channel_index, typ, hs_speed_index)
            print(f"\t\tgot {hs_speed_index=}, {hs_speed=}")

            if hs_speed > best_hs_speed:
                best_ad_channel_index = ad_channel_index
                best_hs_speed_index = hs_speed_index
                best_hs_speed = hs_speed

    ad_channel_index = best_ad_channel_index
    hs_speed_index = best_hs_speed_index
    hs_speed = best_hs_speed

    device.set_ad_channel(ad_channel_index)
    print(f"set {ad_channel_index=}")

    device.set_hs_speed(typ, hs_speed_index)
    print(f"set {typ=}, {hs_speed_index=}, {hs_speed=}")

    if capabilities.ulSetFunctions & AC_SETFUNCTION_BASELINECLAMP:
        baseline_clamp = 1
        device.set_baseline_clamp(baseline_clamp)
        print(f"set {baseline_clamp=}")

    device.set_image(1, 1, 1, xn, 1, yn)

    # %%
    # capture

    device.start_acquisition()
    print("started acquisition")

    print("waiting ...")
    while True:
        time.sleep(1)
        status = device.get_status()
        print(f"got {status=}")

        if status == DRV_IDLE:
            break

    buffer = ctypes.create_string_buffer(xn * yn * 4)
    device.get_acquired_data(xn * yn, buffer)
    print(f"got {buffer=}")

with open(img_file_path, "wb") as fp:
    fp.write(buffer.value)
print(f"saved {img_file_path}")

image = np.frombuffer(buffer, dtype=np.int32).reshape((yn, xn))
plt.imshow(image)
plt.colorbar()
plt.show()
